# -*-coding=utf-8 -*-
import numpy as np
from astropy.units import Ybarn
import math

def computeCorrelation(X, Y):
	# 计算相关度
    xBar = np.mean(X)
    yBar = np.mean(Y)
    SSR = 0
    varX = 0
    varY = 0
    for i in range(0 , len(X)):
        diffXXBar = X[i] - xBar
        diffYYBar = Y[i] - yBar
        SSR += (diffXXBar * diffYYBar)
        varX += diffXXBar**2
        varY += diffYYBar**2
    
    SST = math.sqrt(varX * varY)
    return SSR / SST

# Polynomial Regression
def polyfit(x, y, degree):
	'''
	计算决定系数
	'''
	results = {}

	coeffs = np.polyfit(x, y, degree) # 拟合曲线，求得相关系数 degree几次方的线性回归

	# Polynomial Coefficients
	results['polynomial'] = coeffs.tolist() # 转化为一个列表

	# r-squred
	p = np.poly1d(coeffs) # p相当于拟合曲线方程
	# fit values, and mean
	yhat = p(x) # 传递x自变量值，求出yhat因变量估计值
	ybar = np.sum(y)/len(y) # y均值
	ssreg = np.sum((yhat-ybar)**2) 
	sstot = np.sum((y-ybar)**2)
	results['determination'] = ssreg/sstot

	return results

testX = [1, 3, 8, 7, 9]
testY = [10, 12, 24, 21, 34]

print "r:",computeCorrelation(testX, testY) 
print "r^2:",str(computeCorrelation(testX, testY)**2) # 简单线性回归 决定系数

print polyfit(testX, testY, 1)["determination"]
print polyfit(testX, testY, 1)